#include "udp.h"
#include "kozos.h"
#include "netdef.h"
#include "dns.h"
#include "lib.h"
#include "ip.h"
#include "netlib.h"

typedef struct _udp_recvport {
    udp_receiver r;
    struct _udp_recvport *next;
} udp_recvport;

static kz_msgbox_id_t udp_mbox = MSGBOX_NULL;
static udp_recvport *recvports;
#define UDP_MAXUSE 10
static int udp_usenum;

void udp_init(void)
{
    udp_mbox = kz_getmsgbox();
    recvports = NULL;
    udp_usenum = 0;
}

void handle_udp(uint8 *packet)
{
    ASSERT(udp_mbox != MSGBOX_NULL, "no udp_mbox");
    if (kz_send(udp_mbox, 'h', (char *)packet, 1, NET_MBOXSIZE) == -1)
        kz_kmfree(packet);
}

void udp_send(uint8 *data, uint16 size, uint16 src_port, uint16 dst_port, const uint8 dst_IP[])
{
    uint8 *buf = kz_kmalloc(sizeof(ethernet_header) + sizeof(ip_header) +
                            sizeof(udp_header) + size);
    udp_header *uhdr = (udp_header *)(buf + sizeof(ethernet_header) + sizeof(ip_header));

    memcpy(uhdr + 1, data, size);
    uhdr->udp_length = size + sizeof(*uhdr);
    uhdr->udp_src_port = src_port;
    uhdr->udp_dst_port = dst_port;
    uhdr->udp_checksum = 0;
    // チェックサム計算はsrc_IPセット時

    ip_send(buf, size + sizeof(udp_header), IP_PROTOCOL_UDP, dst_IP, NULL);
}

static int port_is_free(uint16 port)
{
    udp_recvport *p;
    for (p = recvports; p; p = p->next) {
        if (p->r.port == port)
            return 0;
    }
    return 1;
}

static uint16 freeport()
{
    static uint16 last = 65535;
    while (1) {
        if (++last == 0) last = 49152;
        if (port_is_free(last))
            return last;
    }
}

const udp_receiver *udp_use(uint16 port)
{
    int c;
    ASSERT(udp_mbox != MSGBOX_NULL, "no udp_mbox");
    udp_receiver *p = kz_kmalloc(sizeof(udp_recvport));
    p->port = port;
    p->mbox = kz_getmsgbox();
    kz_send(udp_mbox, 'u', (char *)p, 0, 0);
    kz_recv(p->mbox, &c, NULL);
    if (c == -1) {
        kz_freemsgbox(p->mbox);
        kz_kmfree(p);
        return NULL;
    }
    return p;
}

void udp_close(const udp_receiver *receiver)
{
    ASSERT(udp_mbox != MSGBOX_NULL, "no udp_mbox");
    kz_send(udp_mbox, 'c', (char *)receiver, 0, 0);
    while (1) { // receiver->mboxを空にする
        int cmd;
        char *p;
        kz_recv(receiver->mbox, &cmd, &p);
        if (cmd == 'c') break;
        kz_kmfree(p);
    }
    kz_freemsgbox(receiver->mbox);
    kz_kmfree((udp_receiver *)receiver); // udp_useで付けたconstを外す
}

static int check_udp(ip_header *ihdr)
{
    udp_header *uhdr = (udp_header *)(ihdr + 1);
    return uhdr->udp_length >= sizeof(*uhdr) &&
           uhdr->udp_length == ihdr->ip_total_length - sizeof(*ihdr) &&
           (uhdr->udp_checksum == 0 ||
            udp_checksum((uint8 *)uhdr, uhdr->udp_length,
                         ihdr->ip_src_IP, ihdr->ip_dst_IP) == 0xffff);
}

void udp_main(int argc, char *argv)
{
    int cmd;
    char *p;

    while (1) {
        kz_recv(udp_mbox, &cmd, &p);

        if (cmd == 'h') {
            ip_header *ihdr = (ip_header *)(p + sizeof(ethernet_header));
            udp_header *uhdr = (udp_header *)(ihdr + 1);

            if (check_udp(ihdr)) {
                uint16 dport = uhdr->udp_dst_port;
                udp_recvport *rp;
                for (rp = recvports; rp; rp = rp->next) {
                    if (rp->r.port == dport)
                        break;
                }
                if (rp && kz_send(rp->r.mbox, 0, p, 0, NET_MBOXSIZE) != -1)
                    continue;
            }

            kz_kmfree(p);

        } else if (cmd == 'u') {
            udp_recvport *up = (udp_recvport *)p;
            int ret = -1;

            if (udp_usenum >= UDP_MAXUSE)
                goto useend;
            if (up->r.port == 0) {
                up->r.port = freeport();
            } else {
                if (!port_is_free(up->r.port))
                    goto useend;
            }
            up->next = recvports;
            recvports = up;
            udp_usenum++;
            ret = 0;
useend:
            kz_send(up->r.mbox, ret, NULL, 0, 0);

        } else if (cmd == 'c') {
            udp_recvport **r0, *r, *up = (udp_recvport *)p;

            for (r0 = &recvports, r = recvports; r; r0 = &r->next, r = r->next) {
                if (r == up) {
                    *r0 = r->next;
                    udp_usenum--;
                    break;
                }
            }
            ASSERT(r, "can't close port");

            kz_send(up->r.mbox, 'c', NULL, 0, 0);

        } else {
            ASSERT(0, "udp: bad cmd");
        }
    }
}
